import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from pathlib import Path
import streamlit as st
import atexit
import tempfile
from service.image_processor import ImageProcessor
import shutil
import logging
from image_style_trans.style_processor import StyleProcessor

# 添加项目路径
project_root = Path(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(str(project_root))
# 创建临时目录
temp_dir = tempfile.mkdtemp()
def cleanup():
    try:
        shutil.rmtree(temp_dir)
    except Exception as e:
        logging.error(f"清理临时文件时发生错误: {str(e)}")
# 注册清理函数
atexit.register(cleanup)

# 设置页面标题
st.set_page_config(page_title="图像风格转换工具", layout="wide")

# 创建输出目录
output_dir = project_root / "output"
os.makedirs(output_dir, exist_ok=True)

# 创建选项卡
tab1, tab2 = st.tabs(["单张图像处理", "批量处理JSON数据集"])

with tab1:
    # 标题和说明
    st.title("图像风格转换工具")
    st.markdown("""
    这个工具使用可灵AI的图生图API对图像进行风格转换。

    ### 使用方法：
    1. 上传一张图片
    2. 输入风格描述（例如：油画风格，梵高风格，水彩画风格等）
    3. 可选：输入负向提示词（例如：模糊，扭曲，低质量等）
    4. 点击"开始处理"按钮
    5. 等待处理完成后查看结果

    ### 注意事项：
    - 处理过程可能需要一些时间，请耐心等待
    - 图像处理结果会保存在output/styled_images目录下
    - 需要设置KELIN_API_KEY环境变量才能使用此功能
    """)

    st.header("图像风格转换")
    st.markdown("""
    使用可灵AI的图生图API对图像进行风格转换。
    上传一张图片，并输入风格描述，即可生成具有新风格的图像。
    """)
    style_processor = StyleProcessor()
    
    # 检查服务可用性
    if not style_processor.available:
        st.error("图像风格转换服务未正确初始化，请检查API密钥和OSS配置")
        if not style_processor.kelin_service.available:
            st.info("请设置KELIN_API_KEY环境变量")
        if not style_processor.oss_manager.available:
            st.info("请设置OSS_ACCESS_KEY_ID和OSS_ACCESS_KEY_SECRET环境变量")
    
    # 上传图像
    uploaded_file = st.file_uploader("上传图像", type=["jpg", "jpeg", "png"])
    
    # 风格设置
    prompt = st.text_area("风格描述", placeholder="例如：油画风格，梵高风格，水彩画风格等", key="tab1_prompt")
    negative_prompt = st.text_area("负向提示词（可选）", placeholder="例如：模糊，扭曲，低质量等", key="tab1_negative")

    # 处理按钮
    process_button = st.button("开始处理")
    
    # 创建结果占位符
    result_placeholder = st.empty()
    
    if process_button and uploaded_file is not None and prompt:
        # 显示处理中提示
        with st.spinner("正在处理图像，请稍候...这可能需要一些时间"):
            # 保存上传的文件
            temp_file_path = os.path.join(temp_dir, uploaded_file.name)
            with open(temp_file_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
            
            # 创建输出目录
            style_output_dir = output_dir / "styled_images"
            os.makedirs(style_output_dir, exist_ok=True)
            
            # 处理图像
            result_path = style_processor.process_image(
                image_path=temp_file_path,
                prompt=prompt,
                negative_prompt=negative_prompt if negative_prompt else None,
                output_dir=style_output_dir
            )
            
            if result_path:
                # 显示结果
                result_placeholder.success(f"图像处理成功！")
                
                # 创建两列布局
                col1, col2 = st.columns(2)
                
                # 显示原始图像
                with col1:
                    st.subheader("原始图像")
                    st.image(temp_file_path)
                
                # 显示处理后的图像
                with col2:
                    st.subheader("处理后的图像")
                    st.image(result_path)
                
                # 提供下载链接
                with open(result_path, "rb") as f:
                    st.download_button(
                        label="下载处理后的图像",
                        data=f.read(),
                        file_name=os.path.basename(result_path),
                        mime="image/jpeg"
                    )
            else:
                result_placeholder.error("图像处理失败，请检查日志获取详细信息")
    elif process_button:
        if not uploaded_file:
            st.warning("请上传图像")
        if not prompt:
            st.warning("请输入风格描述")

with tab2:
    # 批量处理JSON数据集
    st.title("批量处理JSON数据集")
    st.markdown("""
    这个功能可以批量处理单词数据集JSON文件中的图像，对每个单词的句子图像进行风格转换。

    ### 使用方法：
    1. 选择要处理的JSON文件
    2. 输入风格描述（可选，默认使用"转换为蜡笔彩绘风格"）
    3. 可选：输入负向提示词
    4. 点击"开始批量处理"按钮
    5. 等待处理完成

    ### 注意事项：
    - 处理过程可能需要较长时间，请耐心等待
    - 处理结果会保存在原JSON文件中，并添加styled_path字段
    - 图像将保存在output/images/style_trans/{ip_id}目录下
    """)
    
    # 初始化处理器
    processor = ImageProcessor(project_root, output_dir)
    
    # 列出可用的JSON文件
    json_files = []
    for root, dirs, files in os.walk(output_dir):
        for file in files:
            if file.endswith(".json") and file.startswith("word_dataset_"):
                json_files.append(os.path.join(root, file))
    
    if not json_files:
        st.warning("未找到可处理的JSON文件，请确保output目录下有word_dataset_*.json文件")
    else:
        # 选择JSON文件
        selected_file = st.selectbox("选择要处理的JSON文件", json_files, format_func=lambda x: os.path.basename(x))
        
        # 风格设置
        prompt = st.text_area("风格描述（可选）", placeholder="例如：油画风格，梵高风格，水彩画风格等", key="tab2_prompt")
        negative_prompt = st.text_area("负向提示词（可选）", placeholder="例如：模糊，扭曲，低质量等", key="tab2_negative")
        # 处理按钮
        process_button = st.button("开始批量处理")
        
        if process_button:
            # 显示处理中提示
            with st.spinner("正在批量处理图像，请耐心等待..."):
                # 处理数据集
                result = processor.process_word_dataset(
                    dataset_path=selected_file,
                    prompt=prompt if prompt else None,
                    negative_prompt=negative_prompt if negative_prompt else None
                )
                
                if result:
                    st.success(f"批量处理完成！已更新JSON文件：{os.path.basename(selected_file)}")
                    
                    # 显示处理统计信息
                    total_sentences = sum(len(sentences) for sentences in result.values())
                    processed_count = sum(1 for sentences in result.values() for sentence in sentences if "styled_path" in sentence)
                    
                    st.info(f"共处理 {processed_count}/{total_sentences} 个图像")
                else:
                    st.error("批量处理失败，请检查日志获取详细信息")

# 添加页脚
st.markdown("---")
st.markdown("""
<div style="text-align: center">
    <p>© 2023 图像风格转换工具 | 基于可灵AI图生图API</p>
</div>
""", unsafe_allow_html=True)